package org.finesys.ai.listener;

import lombok.RequiredArgsConstructor;
import org.finesys.ai.entity.AigcEmbedStore;
import org.finesys.ai.entity.AigcKnowledge;
import org.finesys.ai.entity.AigcModel;
import org.finesys.ai.service.AigcEmbedStoreService;
import org.finesys.ai.service.AigcKnowledgeService;
import org.finesys.ai.service.AigcModelService;
import org.finesys.chat.core.base.embedding.factory.EmbeddingStoreFactory;
import org.finesys.chat.core.base.embedding.factory.KnowledgeStoreFactory;
import org.finesys.chat.core.base.factory.ChatModelFactory;
import org.springframework.context.ApplicationListener;
import org.springframework.context.event.ContextRefreshedEvent;
import org.springframework.stereotype.Component;

import java.util.List;

/**
 * 缓存预热监听器
 */
@Component
@RequiredArgsConstructor
public class AigcCachePreheatListener implements ApplicationListener<ContextRefreshedEvent> {

    private final AigcModelService aigcModelService;
    private final AigcEmbedStoreService aigcEmbedStoreService;
    private final AigcKnowledgeService aigcKnowledgeService;
    private final EmbeddingStoreFactory embeddingStoreFactory;
    private final KnowledgeStoreFactory knowledgeStoreFactory;
    private final ChatModelFactory chatModelFactory;

    @Override
    public void onApplicationEvent(ContextRefreshedEvent event) {
        if (event.getApplicationContext().getParent() == null) {
            // 初始化模型
            List<AigcModel> aigcModelList = aigcModelService.findAll();
            chatModelFactory.initModel(aigcModelList);
            // 初始化嵌入向量库
            List<AigcEmbedStore> embedStoreList = aigcEmbedStoreService.findAll();
            embeddingStoreFactory.init(embedStoreList);
            // 初始化知识库
            List<AigcKnowledge> knowledgeList = aigcKnowledgeService.findAll();
            knowledgeStoreFactory.init(knowledgeList);
        }
    }
}
